#include "stdafx.h"
#include "TestGroupHandler.h"

CTestGroupHandler::CTestGroupHandler(EnvDTE::Project * pOriginProject, EnvDTE::Project * pTestProject, LPCSTR lpPrecompileHeader, CAddTracker * pAddTracker) :
	m_spOriginProject(pOriginProject), m_spTestProject(pTestProject), m_strPrecompileHeader(lpPrecompileHeader), m_pAddTracker(pAddTracker)
{
}

CTestGroupHandler::~CTestGroupHandler(void)
{
}

BOOL CTestGroupHandler::MakeTestGroup(LPCWSTR lpTemplate, EnvDTE::CodeElementPtr & spTestGroup)
{
	spTestGroup = NULL;
	
	if (DefineGroupName(lpTemplate))
	{
		if (m_spTestGroupElement != NULL)
		{
			spTestGroup = m_spTestGroupElement;
		} else if (CreateTestGroup())
		{
			int nAttempt = 0;
			while ((FALSE == IsNameUsedInProject(m_bstrGroupName, m_spTestProject, m_spTestGroupElement)) && (nAttempt < 100));
			{
				//after creation it seems there is a need in pause to allow DTE parse new item in code model
				//it would be wonderful to know event or statuc property that could be used as a signal
				Sleep(100);
				++nAttempt;
			} 
			spTestGroup = m_spTestGroupElement;
		}	
	}

	return spTestGroup != NULL;
}

BOOL CTestGroupHandler::IsThereTestGroup(LPCWSTR lpTemplate, EnvDTE::CodeElementPtr & spTestGroup)
{
	m_spTestGroupElement = NULL;
	DefineGroupName(lpTemplate);
	spTestGroup = m_spTestGroupElement;

	return spTestGroup != NULL;
}

void CTestGroupHandler::CreateAndFillTestGroup(CElementTree<CodeTreeItem> & rNode)
{
	ATLASSERT(m_spTestGroupElement == NULL);
	ATLASSERT(m_bstrGroupName.Length() != 0);

	if ((m_bstrGroupName.Length()) && (m_spTestGroupElement == NULL))
	{
		CComBSTR bstrH, bstrCpp;
		VCProjectEngineLibrary::VCProjectPtr spTestProject;
		IDispatchPtr spDispatch;
		CHandle hHFile, hCppFile;

		if (DefineTestGroupFileName(bstrH, bstrCpp) &&
			//add files to project
			SUCCEEDED(m_spTestProject->get_Object(&spDispatch)) &&
			(spTestProject = spDispatch) &&
			SUCCEEDED(spTestProject->AddFile(bstrH, &spDispatch)) &&
			SUCCEEDED(spTestProject->AddFile(bstrCpp, &spDispatch))
			)
		{
			CAtlStringA strClassName(m_bstrGroupName);
			CAtlStringA strH("#pragma once\r\n\r\nclass ");
			strH.Append(strClassName);
			strH.Append("\r\n{\r\n\tDECLARE_UT_FUNC_TABLE\r\npublic:\r\n\t");
			strH.Append(strClassName);
			strH.Append("();\r\n\t~");
			strH.Append(strClassName);
			strH.Append("();\r\n\r\n");

			CAtlStringA strCpp;
			CPathA csHPath(CW2A(bstrH.m_str));
			csHPath.StripPath();

			if (m_strPrecompileHeader.GetLength())
				strCpp.Format("#include \"%s\"\r\n", m_strPrecompileHeader);

			strCpp.Append("#include <utassert.h>\r\n#include \"");
			strCpp.Append(csHPath);
			strCpp.Append("\"\r\n\r\nBEGIN_UT_FUNC_TABLE(");
			strCpp.Append(strClassName);
			strCpp.Append(")\r\n");

			CAtlStringA strFuncTableContent;
			CAtlStringA strCppBodyContent;

			strCppBodyContent.Format("%s::%s()\r\n{\r\n}\r\n\r\n%s::~%s()\r\n{\r\n}\r\n\r\n",
				strClassName, strClassName, strClassName, strClassName);

			if (rNode.GetValue().Checked)
			{
				CAtlStringA strFuncName("Test");

				strH.AppendFormat("\t void %s();\r\n", strFuncName);
				strFuncTableContent.AppendFormat("\tUT_FUNC_ENTRY(%s)\r\n", strFuncName);
				strCppBodyContent.AppendFormat("void %s::%s()\r\n{\r\n\tUT_ASSERT2(FALSE, L\"not implemented\");\r\n}\r\n\r\n",
					strClassName, strFuncName);
			}

			POSITION pos = rNode.GetHeadPosition();
			while(pos)
			{
				const CodeTreeItem & rValue = rNode.GetNext(pos).GetValue();
				if ((rValue.Checked) && (rValue.eKind == EnvDTE::vsCMElementFunction))
				{
					CAtlStringA strFuncName(rValue.Name.GetString());
					strFuncName.Append("Test");

					if (rValue.ElementData.eFunctionKind & EnvDTE::vsCMFunctionConstructor)
						strFuncName = "ConstructorTest";
					else if (rValue.ElementData.eFunctionKind & EnvDTE::vsCMFunctionDestructor)
						strFuncName = "DestructorTest";
					else if (rValue.ElementData.eFunctionKind & EnvDTE::vsCMFunctionOperator)
						strFuncName = "OperatorTest";

					strH.AppendFormat("\t void %s();\r\n", strFuncName);
					strFuncTableContent.AppendFormat("\tUT_FUNC_ENTRY(%s)\r\n", strFuncName);
					strCppBodyContent.AppendFormat("void %s::%s()\r\n{\r\n\tUT_ASSERT2(FALSE, L\"not implemented\");\r\n}\r\n\r\n",
						strClassName, strFuncName);
				}
			}

			strH.Append("};\r\n");
			strCpp.Append(strFuncTableContent);
			strCpp.Append("END_UT_FUNC_TABLE\r\n\r\n");
			strCpp.Append(strCppBodyContent);

			hHFile.Attach(CreateFileW(bstrH, GENERIC_WRITE, 0, NULL, CREATE_NEW, FILE_ATTRIBUTE_NORMAL, NULL));
			hCppFile.Attach(CreateFileW(bstrCpp, GENERIC_WRITE, 0, NULL, CREATE_NEW, FILE_ATTRIBUTE_NORMAL, NULL));

			DWORD dw = 0;
			WriteFile(hHFile, strH.GetString(), strH.GetLength(), &dw, NULL);
			hHFile.Close();
			WriteFile(hCppFile, strCpp.GetString(), strCpp.GetLength(), &dw, NULL);
			hCppFile.Close();

			spTestProject->Save();

			if (m_pAddTracker)
				m_pAddTracker->AddTestGroup(bstrCpp);
		}
	}
}

BOOL CTestGroupHandler::DefineGroupName(LPCWSTR lpTemplate)
{
	m_spTestGroupElement = NULL;
	m_bstrGroupName.Empty();

	CAtlStringW wstrTemplate = lpTemplate;
	
	for (int n = 0; n < wstrTemplate.GetLength(); ++n)
	{
		if (false == iswalnum(wstrTemplate[n]))
			wstrTemplate.SetAt(n, L'_');
	}
	CAtlStringW wstrGroupName = wstrTemplate;

	int nUniqueId = 0;
	do
	{
		EnvDTE::CodeElementPtr spCodeElement;
		if (FALSE == IsNameUsedInProject(wstrGroupName, m_spOriginProject, spCodeElement))
		{
			if (IsNameUsedInProject(wstrGroupName, m_spTestProject, spCodeElement) &&
				(spCodeElement))
			{
				if (IsElementTestGroup(spCodeElement))
				{
					m_bstrGroupName = wstrGroupName;
					m_spTestGroupElement = spCodeElement;
					return TRUE;
				}
			} else {
				m_bstrGroupName = wstrGroupName;
				return TRUE; //suitable name have been found
			}
		}
		//try next
		++nUniqueId;
		wstrGroupName = wstrTemplate;
		wstrGroupName.AppendFormat(L"%i", nUniqueId);
	}
	while(nUniqueId < 10000);

	return FALSE;
}

BOOL CTestGroupHandler::IsNameUsedInProject(LPCWSTR lpName, EnvDTE::Project * pProject, EnvDTE::CodeElementPtr & spMatched)
{
	EnvDTE::CodeModelPtr spCodeModel;
	VCCodeModelLibrary::VCCodeModelPtr spVCCodeModel;
	EnvDTE::CodeElementsPtr spCodeElements;
	VCCodeModelLibrary::VCCodeElementsPtr spVCCodeElements;
	IDispatchPtr spDispatch;
	
	if (SUCCEEDED(pProject->get_CodeModel(&spCodeModel)) &&
		(spVCCodeModel = spCodeModel))
	{
		if (SUCCEEDED(spVCCodeModel->get_Classes(&spCodeElements)) &&
			(spVCCodeElements = spCodeElements) &&
			SUCCEEDED(spVCCodeElements->Find(CComBSTR(lpName), &spDispatch)) &&
			(spMatched = spDispatch))
		{
			return TRUE;
		}
		
		if (SUCCEEDED(spVCCodeModel->get_Structs(&spCodeElements)) &&
			(spVCCodeElements = spCodeElements) &&
			SUCCEEDED(spVCCodeElements->Find(CComBSTR(lpName), &spDispatch)) &&
			(spMatched = spDispatch))
		{
			return TRUE;
		}

		if (SUCCEEDED(spVCCodeModel->get_Enums(&spCodeElements)) &&
			(spVCCodeElements = spCodeElements) &&
			SUCCEEDED(spVCCodeElements->Find(CComBSTR(lpName), &spDispatch)) &&
			(spDispatch))
		{
			return TRUE;
		}

		if (SUCCEEDED(spVCCodeModel->get_Functions(&spCodeElements)) &&
			(spVCCodeElements = spCodeElements) &&
			SUCCEEDED(spVCCodeElements->Find(CComBSTR(lpName), &spDispatch)) &&
			(spDispatch))
		{
			return TRUE;
		}
		
		if (SUCCEEDED(spVCCodeModel->get_Typedefs(&spCodeElements)) &&
			(spVCCodeElements = spCodeElements) &&
			SUCCEEDED(spVCCodeElements->Find(CComBSTR(lpName), &spDispatch)) &&
			(spDispatch))
		{
			return TRUE;
		}

		if (SUCCEEDED(spVCCodeModel->get_Unions(&spCodeElements)) &&
			(spVCCodeElements = spCodeElements) &&
			SUCCEEDED(spVCCodeElements->Find(CComBSTR(lpName), &spDispatch)) &&
			(spDispatch))
		{
			return TRUE;
		}

		if (SUCCEEDED(spVCCodeModel->get_Variables(&spCodeElements)) &&
			(spVCCodeElements = spCodeElements) &&
			SUCCEEDED(spVCCodeElements->Find(CComBSTR(lpName), &spDispatch)) &&
			(spDispatch))
		{
			return TRUE;
		}

	}
	return FALSE;
}

BOOL CTestGroupHandler::CreateTestGroup()
{
	//Seems it is much easier to add new files to a project than create code in CodeModel
	CComBSTR bstrH, bstrCpp;
	VCProjectEngineLibrary::VCProjectPtr spTestProject;
	IDispatchPtr spDispatch;

	if (DefineTestGroupFileName(bstrH, bstrCpp) &&
		//add files to project
		SUCCEEDED(m_spTestProject->get_Object(&spDispatch)) &&
		(spTestProject = spDispatch) &&
		SUCCEEDED(spTestProject->AddFile(bstrH, &spDispatch)) &&
		SUCCEEDED(spTestProject->AddFile(bstrCpp, &spDispatch)) &&
		//fill content of H file
		CreateTestGroupHFile(bstrH) &&
		//fill contetn of Cpp file
		CreateTestGroupCppFile(bstrCpp, bstrH) &&
		CreateTestGroupClass(bstrH, bstrCpp) &&
		SUCCEEDED(spTestProject->Save()))
	{
		//find new test class in code model
		return TRUE;
	}
	
	return FALSE;
}

BOOL CTestGroupHandler::DefineTestGroupFileName(CComBSTR & bstrH, CComBSTR & bstrCpp)
{
	BOOL bRet = FALSE;
	CComBSTR bstrProjPath;
	if ((m_bstrGroupName.Length()) &&
		SUCCEEDED(m_spTestProject->get_FileName(&bstrProjPath)))
	{
		CPathW csProjPath = bstrProjPath;
		csProjPath.UnquoteSpaces();
		csProjPath.RemoveFileSpec();

		LPCWSTR lpTemplate = m_bstrGroupName[0] == L'C' ? m_bstrGroupName.m_str + 1 : m_bstrGroupName.m_str;
		csProjPath.Append(lpTemplate);
		CAtlStringW wstrHFile(csProjPath);
		CAtlStringW wstrCppFile(csProjPath);
		int nUniqueId = 0;

		do
		{
			wstrHFile.Append(L".h");
			wstrCppFile.Append(L".cpp");
		
			if ((INVALID_FILE_ATTRIBUTES == GetFileAttributesW(wstrHFile)) &&
				(INVALID_FILE_ATTRIBUTES == GetFileAttributesW(wstrCppFile)))
			{
				//bingo - both file names are unique for the test project directory
				bstrH = wstrHFile;
				bstrCpp = wstrCppFile;
				bRet = TRUE;
				break;
			}

			++nUniqueId;
			wstrHFile = csProjPath.m_strPath;
			wstrHFile.AppendFormat(L"%i", nUniqueId);
			wstrCppFile = csProjPath.m_strPath;
			wstrCppFile.AppendFormat(L"%i", nUniqueId);
		}
		while ( nUniqueId < 10000);
	}

	return bRet;
}

BOOL CTestGroupHandler::CreateTestGroupHFile(LPCWSTR lpPath)
{
	
	CAtlStringA strContent = "#pragma once\r\n\r\n";

	BOOL bRet = FALSE;
	HANDLE hFile = CreateFileW(lpPath, GENERIC_WRITE, 0, NULL, CREATE_NEW, FILE_ATTRIBUTE_NORMAL, NULL);
	if (INVALID_HANDLE_VALUE != hFile)
	{
		DWORD dw = 0;
		BOOL bWrite = WriteFile(hFile, strContent.GetString(), strContent.GetLength(), &dw, NULL);
		BOOL bClose = CloseHandle(hFile);
		bRet = bWrite && bClose;
	}

	return bRet;
}

BOOL CTestGroupHandler::CreateTestGroupCppFile(LPCWSTR lpPath, LPCWSTR lpHFilePath)
{
	CAtlStringA strGroupName(m_bstrGroupName);
	CAtlStringA strHFilePath(lpHFilePath);
	CPathA csHPath(strHFilePath);
	csHPath.StripPath();
	CAtlStringA strContent;
	if (m_strPrecompileHeader.GetLength())
	{
		strContent.Append("#include \"");
		strContent.Append(m_strPrecompileHeader);
		strContent.Append("\"\r\n");
	}

	strContent.Append("#include <utassert.h>\r\n#include \"");
	strContent.Append(csHPath);
	strContent.Append("\"\r\n\r\nBEGIN_UT_FUNC_TABLE(");
	strContent.Append(strGroupName);
	strContent.Append(")\r\nEND_UT_FUNC_TABLE\r\n\r\n\r\n");

	BOOL bRet = FALSE;
	HANDLE hFile = CreateFileW(lpPath, GENERIC_WRITE, 0, NULL, CREATE_NEW, FILE_ATTRIBUTE_NORMAL, NULL);
	if (INVALID_HANDLE_VALUE != hFile)
	{
		DWORD dw = 0;
		BOOL bWrite = WriteFile(hFile, strContent.GetString(), strContent.GetLength(), &dw, NULL);
		BOOL bClose = CloseHandle(hFile);
		bRet = bWrite && bClose;
	}

	return bRet;
}

BOOL CTestGroupHandler::CreateTestGroupClass(LPCWSTR lpHFilePath, LPCWSTR lpCppFilePath)
{
	//REMARK
	//It's easier to form the test group class during creation source file creation
	//but it won't be present in the studio internal model that makes futher operations impossible
	BOOL bRet = FALSE;

	EnvDTE::CodeModelPtr spCodeModel;
	VCCodeModelLibrary::VCCodeModelPtr spVCCodeModel;
	EnvDTE::CodeClassPtr spCodeClass;
	EnvDTE::CodeFunctionPtr spCodeFunction;

	EnvDTE::TextPointPtr spStart, spEnd;
	EnvDTE::EditPointPtr spEditPoint;
	CComBSTR bstrBody;
	
	if (SUCCEEDED(m_spTestProject->get_CodeModel(&spCodeModel)) && 
		(spVCCodeModel = spCodeModel) &&
		SUCCEEDED(spVCCodeModel->AddClass(m_bstrGroupName, CComVariant(lpHFilePath), CComVariant(0), CComVariant(0), CComVariant(0), EnvDTE::vsCMAccessPublic, &spCodeClass)) &&
		(spCodeClass) &&
		SUCCEEDED(spCodeClass->get_StartPoint(&spStart)) && (spStart) &&
		SUCCEEDED(spCodeClass->get_EndPoint(&spEnd)) && (spEnd) &&
		SUCCEEDED(spStart->CreateEditPoint(&spEditPoint)) && (spEditPoint) &&
		SUCCEEDED(spEditPoint->GetText(CComVariant((LPDISPATCH)spEnd), &bstrBody)))
	{
		CAtlStringW wstrBody = bstrBody;
		wstrBody.Remove(L'\r');
		long nOpenBracket = wstrBody.Find(L'{');

		if (SUCCEEDED(spEditPoint->CharRight(nOpenBracket + 1)) &&
			SUCCEEDED(spEditPoint->Insert(CComBSTR(L"\r\nDECLARE_UT_FUNC_TABLE\r\n"))) &&
			SUCCEEDED(spCodeClass->AddFunction(m_bstrGroupName, EnvDTE::vsCMFunctionConstructor, CComVariant(EnvDTE::vsCMTypeRefOther), CComVariant(), EnvDTE::vsCMAccessPublic, CComVariant(lpCppFilePath), &spCodeFunction))
			)
		{
			bRet = TRUE;
		}
	}
	return bRet;
}